import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import cooler
import h5py
import hicstraw
import argparse
from scipy import sparse

def load_hic_file(filename):
    """
    Load a .hic file using hicstraw
    
    Parameters:
    -----------
    filename : str
        Path to .hic file
    
    Returns:
    --------
    hic : hicstraw.HiCFile
        HiCFile object
    """
    print(f"Loading HiC file: {filename}")
    return hicstraw.HiCFile(filename)

def load_cooler_file(filename):
    """
    Load a .cool file using cooler
    
    Parameters:
    -----------
    filename : str
        Path to .cool file
    
    Returns:
    --------
    cool : cooler.Cooler
        Cooler object
    """
    print(f"Loading Cooler file: {filename}")
    return cooler.Cooler(filename)

def get_hic_matrix(hic, chromosome1, chromosome2=None, resolution=10000, normalization="KR"):
    """
    Extract contact matrix from .hic file
    
    Parameters:
    -----------
    hic : hicstraw.HiCFile
        HiCFile object
    chromosome1 : str
        First chromosome (e.g., "chr1")
    chromosome2 : str, optional
        Second chromosome (e.g., "chr2"). If None, uses chromosome1
    resolution : int, optional
        Resolution in base pairs
    normalization : str, optional
        Normalization method, one of "NONE", "VC", "VC_SQRT", "KR"
        
    Returns:
    --------
    matrix : numpy.ndarray
        Contact matrix
    """
    if chromosome2 is None:
        chromosome2 = chromosome1
    
    print(f"Extracting matrix for {chromosome1}-{chromosome2} at {resolution}bp resolution using {normalization} normalization")
    
    # Get matrix from HiC file
    matrix_obj = hic.getMatrixZoomData(chromosome1, chromosome2, "observed", normalization, "BP", resolution)
    
    # Get dimensions
    bins1 = matrix_obj.getBinNumberAtPosition(chromosome1, 0, matrix_obj.chr1Length)
    bins2 = matrix_obj.getBinNumberAtPosition(chromosome2, 0, matrix_obj.chr2Length)
    
    # Create empty matrix
    matrix = np.zeros((bins1, bins2))
    
    # Fill matrix with data
    records = matrix_obj.getRecords()
    for record in records:
        bin1 = record.binX
        bin2 = record.binY
        count = record.counts
        matrix[bin1, bin2] = count
        if chromosome1 == chromosome2 and bin1 != bin2:
            matrix[bin2, bin1] = count
    
    return matrix

def get_cooler_matrix(cool, region1, region2=None, balance=True):
    """
    Extract contact matrix from cooler file
    
    Parameters:
    -----------
    cool : cooler.Cooler
        Cooler object
    region1 : str
        First genomic region (e.g., "chr1:1000000-2000000")
    region2 : str, optional
        Second genomic region. If None, uses region1
    balance : bool, optional
        Whether to balance the matrix
        
    Returns:
    --------
    matrix : numpy.ndarray
        Contact matrix
    """
    if region2 is None:
        region2 = region1
    
    print(f"Extracting matrix for {region1}-{region2}, balanced: {balance}")
    
    # Get matrix from cooler file
    matrix = cool.matrix(balance=balance).fetch(region1, region2)
    
    return matrix

def visualize_matrix(matrix, title="Contact Matrix", cmap="YlOrRd", log_scale=True, vmin=None, vmax=None, output=None):
    """
    Visualize contact matrix
    
    Parameters:
    -----------
    matrix : numpy.ndarray
        Contact matrix
    title : str, optional
        Plot title
    cmap : str, optional
        Matplotlib colormap
    log_scale : bool, optional
        Whether to log-transform the matrix
    vmin : float, optional
        Minimum value for color scale
    vmax : float, optional
        Maximum value for color scale
    output : str, optional
        Path to save figure. If None, displays the figure
    """
    # Create a copy of the matrix for visualization
    viz_matrix = matrix.copy()
    
    # Handle zeros for log scale
    if log_scale:
        viz_matrix = np.log1p(viz_matrix)
        title += " (log scale)"
    
    # Create figure
    plt.figure(figsize=(10, 10))
    plt.imshow(viz_matrix, cmap=cmap, vmin=vmin, vmax=vmax, interpolation="nearest")
    plt.colorbar(label="Contact frequency")
    plt.title(title)
    plt.xlabel("Genomic position")
    plt.ylabel("Genomic position")
    
    # Save or display figure
    if output:
        plt.savefig(output, dpi=300, bbox_inches="tight")
        print(f"Figure saved to {output}")
    else:
        plt.show()

def main():
    """Main function to process command line arguments"""
    parser = argparse.ArgumentParser(description="Process and visualize Hi-C data")
    parser.add_argument("--file", "-f", type=str, required=True, help="Path to Hi-C file (.hic or .cool format)")
    parser.add_argument("--format", "-t", type=str, choices=["hic", "cool", "auto"], default="auto", 
                       help="Format of input file ('hic', 'cool', or 'auto' to guess from extension)")
    parser.add_argument("--region1", "-r1", type=str, required=True, help="First region (e.g., 'chr1' or 'chr1:1000000-2000000')")
    parser.add_argument("--region2", "-r2", type=str, help="Second region (default: same as first region)")
    parser.add_argument("--resolution", "-res", type=int, default=10000, help="Resolution in base pairs for .hic files")
    parser.add_argument("--normalization", "-norm", type=str, default="KR",
                       choices=["NONE", "VC", "VC_SQRT", "KR"], help="Normalization method for .hic files")
    parser.add_argument("--balance", "-b", action="store_true", help="Balance matrix for .cool files")
    parser.add_argument("--output", "-o", type=str, help="Path to save visualization (if not provided, displays figure)")
    parser.add_argument("--log", "-l", action="store_true", help="Use log scale for visualization")
    parser.add_argument("--cmap", "-c", type=str, default="YlOrRd", help="Matplotlib colormap for visualization")
    parser.add_argument("--vmin", type=float, help="Minimum value for color scale")
    parser.add_argument("--vmax", type=float, help="Maximum value for color scale")
    
    args = parser.parse_args()
    
    # Determine file format
    file_format = args.format
    if file_format == "auto":
        if args.file.endswith(".hic"):
            file_format = "hic"
        elif args.file.endswith(".cool") or args.file.endswith(".mcool"):
            file_format = "cool"
        else:
            raise ValueError("Could not determine file format from extension. Please specify --format.")
    
    # Load file
    if file_format == "hic":
        hic = load_hic_file(args.file)
        matrix = get_hic_matrix(hic, args.region1, args.region2, args.resolution, args.normalization)
    else:  # cool
        cool = load_cooler_file(args.file)
        matrix = get_cooler_matrix(cool, args.region1, args.region2, args.balance)
    
    # Visualize matrix
    title = f"Contact Matrix: {args.region1}"
    if args.region2 and args.region2 != args.region1:
        title += f" vs {args.region2}"
    
    visualize_matrix(matrix, title=title, cmap=args.cmap, log_scale=args.log,
                    vmin=args.vmin, vmax=args.vmax, output=args.output)

if __name__ == "__main__":
    main()
